package com.deepface.web.controller.faceLab;

/**
 * @ClassName WebSocketController
 * @Description
 * @Author 74707
 * @Date 2024/4/1 16:12
 * @Version V1.0
 */

import com.deepface.common.utils.StringUtils;
import com.deepface.framework.config.MyEndpointConfigure;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.CrossOrigin;

/**
 * WebSocket获取实时日志并输出到Web页面
 */
@CrossOrigin
@Component
@ServerEndpoint(value = "/websocket", configurator = MyEndpointConfigure.class)
public class WebSocketController {

  private static final ConcurrentHashMap<String, Session> sessionMap = new ConcurrentHashMap<>();
  private static final Map<String, Set<Session>> jobIdSessionMap = new ConcurrentHashMap<>();


  // 添加Session到jobId的映射中
  private void addSessionToJobIdMapping(String jobId, Session session) {
    jobIdSessionMap.computeIfAbsent(jobId, k -> new HashSet<>()).add(session);
  }

  // 从jobId的映射中移除Session
  private void removeSessionFromJobIdMapping(String jobId, Session session) {
    Set<Session> sessions = jobIdSessionMap.get(jobId);
    if (sessions != null) {
      sessions.remove(session);
      if (sessions.isEmpty()) {
        jobIdSessionMap.remove(jobId);
      }
    }
  }

  @OnOpen
  public void onOpen(Session session) {
    // 将新会话添加到map中
    sessionMap.put(session.getId(), session);
    System.out.println("New connection: " + session.getId());
  }

  @OnClose
  public void onClose(Session session) {
    sessionMap.remove(session.getId());
    // 遍历 jobIdSessionMap 移除该 session
    jobIdSessionMap.forEach((jobId, sessions) -> {
      if (sessions.contains(session)) {
        removeSessionFromJobIdMapping(jobId, session);
      }
    });
    System.out.println("Connection closed: " + session.getId());
  }

  @OnMessage
  public void onMessage(String message, Session session) {
    System.out.println("------我手打-----"+message);
    if (message.startsWith("subscribe:")) {
      String jobId = message.substring("subscribe:".length());
      addSessionToJobIdMapping(jobId, session);
      System.out.println("Session " + session.getId() + " subscribed to jobId " + jobId);
    } else {
      // 处理其他类型的消息，例如日志消息
      String[] parts = message.split(":", 2);
      if (parts.length != 2) {
        System.out.println("Invalid message format: " + message);
        return;
      }
      String jobId = parts[0];
      String line = parts[1];
      System.out.println(line);
      sendToSpecificClient(jobId, line);
    }

  }


  // 发送消息到指定客户端
  public static void sendToSpecificClient(String jobId, String line) {
    Set<Session> sessions = jobIdSessionMap.get(jobId);
    if (sessions != null) {
      for (Session session : sessions) {
        if (session.isOpen()) {
          try{
            session.getBasicRemote().sendText(line);
          }catch (Exception e){
            e.printStackTrace();
          }
        }
      }
    }
}

  // 广播消息到所有客户端
  public static void broadcast(String message) {
    sessionMap.forEach((id, session) -> {
      if (session.isOpen()) {
        try {
          session.getBasicRemote().sendText(message);
        } catch (IOException e) {
          e.printStackTrace();
        }
      }
    });
  }

}
